{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses #-}

module Common where

import Prelude  hiding ( concat, foldl, foldr, mapM, foldl1 )
import Debug.Trace
import Control.Applicative
import Control.Arrow ( (>>>), (***), first, second )
import Control.Monad.Identity hiding ( mapM, forM )
import Control.Monad.Writer hiding ( mapM, forM )
import Control.Monad.Reader hiding ( mapM, forM )
import Control.Monad.State hiding ( mapM, forM )
import Control.Monad.ST
import Control.DeepSeq
import Data.Array.ST
import Data.STRef
import Data.Monoid
import Data.Char
import Data.Maybe
import Data.Traversable
import Data.Foldable 
import qualified Data.List as List
import Data.Map ( Map )
import Data.IntMap ( IntMap )
import qualified Data.Map as Map
import qualified Data.IntMap as IMap
import System.Random

class IdCounter a where
  newId :: a -> (Int, a)
  
($$!) :: NFData a => (a -> b) -> a -> b
f $$! x = rnf x `seq` f x
  
newIdS :: IdCounter a => State a Int
newIdS = State newId
  
newtype EndoKleisli m a b = 
  EndoKleisli { runEndoKleisli :: (a -> m a, b) }

instance (Monad m, Monoid b) => Monoid (EndoKleisli m a b) where
  mempty = EndoKleisli (return, mempty)
  (EndoKleisli (f, a)) `mappend` (EndoKleisli (g, b)) = 
    EndoKleisli (f >=> g, a `mappend` b)
  
instance Monad m => Monad (EndoKleisli m a) where
  return x = EndoKleisli (return, x)
  (EndoKleisli (f, a)) >>= g = 
    let EndoKleisli (h, b) = g a in EndoKleisli (f >=> h, b)

mapKeysM :: (Ord k2, Monad m) => (k1 -> m k2) -> Map k1 v -> m (Map k2 v)
mapKeysM f m = do
  let mapFirstM (a, b) = do
        a' <- f a
        return (a', b)
  m' <- mapM mapFirstM (Map.toList m)
  return (Map.fromList m')
  
mapPair :: (a -> b) -> (a, a) -> (b, b)
mapPair f = f *** f
  
mapPairM :: Monad m => (a -> m b) -> (a, a) -> m (b, b)
mapPairM f (x, y) = do
  x' <- f x
  y' <- f y
  return (x', y')
    
mapM_to_fmap :: 
  ((a -> Identity b) -> c -> Identity d) -> (a -> b) -> c -> d
mapM_to_fmap mf f = mf (f >>> Identity) >>> runIdentity

mapM_to_foldMap :: (Monoid m) => 
  ((a -> Writer m a) -> b -> Writer m c) -> (a -> m) -> b -> m
mapM_to_foldMap mf f = mf f' >>> execWriter
  where f' x = tell (f x) >> return x

mapM_to_foldMapA :: (Monoid m, Applicative t) => 
  ((a -> Writer m (t a)) -> b -> Writer m c) -> (a -> m) -> b -> m
mapM_to_foldMapA mf f = mf f' >>> execWriter
  where f' x = tell (f x) >> return (pure x)
  
mapM_to_foldMapAM :: (Monoid m, Monad f, Applicative t) =>
  ((a -> WriterT m f (t a)) -> b -> WriterT m f c) -> (a -> f m) -> b -> f m
mapM_to_foldMapAM mf f = mf (f_gen f) >>> execWriterT
  where 
    f_gen :: (Monoid m, Monad f, Applicative t) => (a -> f m) -> a -> WriterT m f (t a)
    f_gen f x = do
      x' <- lift $ f x
      tell x'
      return (pure x)
      
mapM_to_foldMapM :: (Monoid m, Monad f) =>
  ((a -> WriterT m f a) -> b -> WriterT m f c) -> (a -> f m) -> b -> f m
mapM_to_foldMapM mf f = mf (f_gen f) >>> execWriterT
  where 
    f_gen :: (Monoid m, Monad f) => 
      (a -> f m) -> a -> WriterT m f a
    f_gen f x = do
      x' <- lift $ f x
      tell x'
      return x
      
foldMap_to_foldr :: ((a -> Endo b) -> c -> Endo b) -> 
  (a -> b -> b) -> b -> c -> b
foldMap_to_foldr foldMap f z t = appEndo (foldMap (Endo . f) t) z

mapM_to_foldrM :: (Monad m, Applicative t) => 
  ((a -> EndoKleisli m b (t a)) -> c -> EndoKleisli m b d) -> 
    (a -> b -> m b) -> b -> c -> m b
mapM_to_foldrM mapM f z t = 
  (fst $ runEndoKleisli $ mapM (fld f) t) z
  where
    fld :: (Monad m, Applicative t) => 
      (a -> b -> m b) -> a -> EndoKleisli m b (t a)
    fld f x = EndoKleisli (f x, pure x)
    
flipPair :: (a, b) -> (b, a)
flipPair (a, b) = (b, a)

anyM :: (Monad f, Traversable t) => (a -> f Bool) -> t a -> f Bool
anyM f = mapM f >>> liftM (foldr (||) False)

allM :: (Monad f, Traversable t) => (a -> f Bool) -> t a -> f Bool
allM f = mapM f >>> liftM (foldr (&&) True)
  
infixr 9 <.>
(<.>) :: Functor f => (b -> c) -> (a -> f b) -> (a -> f c) 
(<.>) = ((.) . fmap)

replace :: (Functor f, Eq a) => a -> a -> f a -> f a
replace = genericReplace fmap

replaceMany :: (Foldable t, Functor f, Eq a) => t (a, a) -> f a -> f a
replaceMany = flip $ foldr (uncurry replace) 

genericReplace :: Eq a => ((a -> a) -> b -> b) -> a -> a -> b -> b
genericReplace f x y = f $ \a -> if a == x then y else a

genericReplaceMany :: (Foldable t, Eq a) => 
  ((a -> a) -> b -> b) -> t (a, a) -> b -> b
genericReplaceMany f = flip $ foldr (uncurry (genericReplace f))

concatMapM :: Monad m => (a -> m [b]) -> [a] -> m [b]
concatMapM f = (mapM f) >>> (liftM concat)

mapMaybeM :: Monad m => (a -> m (Maybe b)) -> [a] -> m [b]
mapMaybeM f = (mapM f) >>> (liftM (filter isJust >>> map fromJust))

containsDuplicates :: (Foldable t, Eq a) => t a -> Bool
containsDuplicates = toList >>> (\x -> length (List.nub x) /= length x)

appendEntry :: (Ord k, Monoid a) => k -> a -> Map k a -> Map k a
appendEntry k a = Map.alter f k
  where f Nothing = Just a
        f (Just a') = Just $ a' `mappend` a

type ReducedT = WriterT Any

traceOutput :: Show a => a -> a
traceOutput x = trace (show x) x

reduced :: Monad m => ReducedT m ()
reduced = tell (Any True)

reduce :: Monad m => [(a -> ReducedT m a)] -> a -> m a
reduce fs = 
  reduceOne (foldl1 (>=>) (map reduceOne fs)) >>> 
  runWriterT >>> (liftM fst)
  where
    reduceOne :: Monad m => (a -> ReducedT m a) -> a -> ReducedT m a
    reduceOne f a = do
      (a', r) <- listen $ f a
      if getAny r then reduceOne f a' else return a'
      
clusterBy :: (a -> a -> Bool) -> [a] -> [[a]]
clusterBy _ [] = []
clusterBy f (a : as) = (a : as') : (clusterBy f as'')
  where (as', as'') = List.partition (f a) as

showId :: Int -> String
showId = ('?' :) . intToChars
  
intToChars :: Int -> String
intToChars 0 = []
intToChars n = 
  let c = chr $ (n `mod` 26) + (ord 'a')
  in c : intToChars (n `div` 26)  
  
newtype UnquotedString = Unquoted String
  deriving ( Eq, Ord )
  
instance Show UnquotedString where
  show (Unquoted s) = s
  
showUnquoted :: (Functor t, Show (t UnquotedString), 
  Show (t String)) => t String -> String
showUnquoted = fmap Unquoted >>> show

readonly :: Reader s a -> State s a
readonly r = State $ \s -> (runReader r s, s)

stateful :: State r a -> (a -> Reader r b) -> Reader r b
stateful st f = do
  (a, r) <- runState st <$> ask
  local (const r) (f a)
  
statelocal :: State r b -> Reader r a -> Reader r a
statelocal st = local $ \r -> execState st r 

type Indented = Reader Int

indentation :: Indented String
indentation = do
  i <- ask
  return $ '\n' : (concat $ replicate i "  ")
  
indent :: Indented a -> Indented a
indent = local (+ 1)

runIndented :: Indented a -> a
runIndented = flip runReader 0

instance Foldable ((,) a) where
  foldMap f (_, x) = f x
  
instance Traversable ((,) a) where
  sequenceA (s, x) = ((,) s) <$> x
  
nubbed x = length (List.nub x) == length x 

shuffle :: [a] -> StdGen -> ([a], StdGen)
shuffle xs gen = runST $ do
  g <- newSTRef gen
  let randomRST lohi = do
        (a,s') <- liftM (randomR lohi) (readSTRef g)
        writeSTRef g s'
        return a
  ar <- newArray n xs
  xs' <- forM [1..n] $ \i -> do
          j <- randomRST (i,n)
          vi <- readArray ar i
          vj <- readArray ar j
          writeArray ar j vi
          return vj
  gen' <- readSTRef g
  return (xs', gen')
  where
    n = length xs
    newArray :: Int -> [a] -> ST s (STArray s Int a)
    newArray n xs =  newListArray (1,n) xs
    
class JSON a where
  json :: a -> String

